import os
from langchain_openai import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.agents import initialize_agent
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferWindowMemory,ConversationBufferMemory    



llm = ChatOpenAI(
    temperature=0.05 ,
    model="glm-4-0520",
    openai_api_base="https://open.bigmodel.cn/api/paas/v4/",
    streaming=True,
)




from langchain.chains import LLMMathChain
from langchain.agents import Tool
 
llm_math = LLMMathChain(llm = llm)


# llm_math.invoke("计算1*10**2+3*0")
# exit()
 
# initialize the math tool
math_tool = Tool(
    name ='Calculator',
    func = llm_math.run,
    description ='对于解决数学和计算问题非常有效的工具。'
)




from langchain_community.tools.tavily_search import TavilySearchResults


tools = [
    math_tool,
    TavilySearchResults(
        search_result=True,
        max_results=3
    ),
]


# 获取要使用的提示模板 - 您可以修改此模板！
from langchain import hub
# # prompt = hub.pull("hwchase17/react-chat-json")


agent_kwargs = {
    "extra_prompt_messages": [MessagesPlaceholder(variable_name="chat_history")],
}
memory = ConversationBufferWindowMemory(
    memory_key="chat_history",
    # return_messages=True,
    k=3,
)


from langchain.agents import AgentExecutor, create_json_chat_agent
from langchain.agents import AgentType

prompt = ChatPromptTemplate(
    messages=[
        # agent是一个理财小助手，agent的
    ]
)

agent = initialize_agent(
    llm=llm,
    verbose=True,
    tools=tools,
    # 聊天机器人的类型
    agent_type=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
    agent_kwargs=agent_kwargs,
    memory=memory
)


agent.run("I'm Mark")
# agent.run("Are you happy?")
agent.run("What's my name?")


def Get_agent() -> AgentExecutor:
    return agent
